from celery import shared_task
from zai import ZhipuAiClient
import time
import logging
import os
from dotenv import load_dotenv
from .models import VideoHistory
from user.models import SysUser

logger = logging.getLogger(__name__)
load_dotenv()

@shared_task
def generate_video(user_id, prompt, mode, image_url=None, confirmed=False, quality="quality", fps=30):
    try:
        if not confirmed:
            logger.info(f"用户 {user_id} 的任务，提示词 '{prompt}' 等待确认")
            return {'status': 'WAITING', 'prompt': prompt, 'image_url': image_url, 'mode': mode, 'quality': quality, 'fps': fps}

        api_key = os.getenv('ZHIPUAI_API_KEY')
        if not api_key:
            raise ValueError("未设置 ZHIPUAI_API_KEY 环境变量")
        client = ZhipuAiClient(api_key=api_key)
        response = client.videos.generations(
            model="cogvideox-3",
            prompt=prompt,
            image_url=image_url,
            quality=quality,
            with_audio=True,
            size="1920x1080",
            fps=fps,
        )
        logger.info(f"初始响应: {response}")

        task_id = response.id
        max_attempts = 120  # 20 分钟
        for _ in range(max_attempts):
            status = client.videos.retrieve_videos_result(id=task_id)
            logger.info(f"任务状态: {status.task_status}")
            if status.task_status == 'SUCCESS':
                if status.video_result and len(status.video_result) > 0:
                    video_url = status.video_result[0].url
                    logger.info(f"视频结果: {video_url}")
                    # 保存到 VideoHistory
                    user = SysUser.objects.get(id=user_id)
                    image_data = image_url if mode == 'image' else (','.join(image_url) if mode == 'frames' and image_url else '')
                    history = VideoHistory.objects.create(
                        user=user,
                        task_id=task_id,
                        prompt=prompt,
                        mode=mode,
                        image_url=image_data,
                        video_url=video_url,  # 直接使用远程URL
                        status='SUCCESS',
                        quality=quality,
                        fps=fps
                    )
                    return {'url': video_url}
                raise Exception("未找到视频结果")
            elif status.task_status == 'FAILURE':
                raise Exception(f"视频生成失败: {status.error_message or '未知错误'}")
            logger.info("任务仍在处理中，等待...")
            time.sleep(10)
        raise Exception("任务在20分钟后超时")
    except Exception as e:
        logger.error(f"视频生成错误: {str(e)}")
        raise